import subprocess
import os
import time
import requests
import shutil
from datetime import datetime
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler

class MCPServerManager:
    def __init__(self, main_script_path, watch_directory):
        self.main_script_path = main_script_path
        self.watch_directory = watch_directory
        self.process = None
        self.observer = Observer()
        self.event_handler = self._create_event_handler()
        self.last_restart_time = 0
        self.restart_debounce_seconds = 3  # Minimum time between restarts
        self.health_check_url = "http://localhost:8000/health"
        self.backup_dir = os.path.join(os.path.dirname(main_script_path), "backups")
        self.last_known_good_script = None
        self.startup_timeout = 30  # seconds to wait for server startup
        self.health_check_retries = 3
        # function_registry.json路径（用于备份）
        self.registry_path = os.path.join(os.path.dirname(main_script_path), "function_registry.json")
        os.makedirs(self.backup_dir, exist_ok=True)

    def _create_event_handler(self):
        class Handler(FileSystemEventHandler):
            def __init__(self, manager):
                self.manager = manager

            def on_any_event(self, event):
                if event.is_directory or not event.src_path.endswith('.py'):
                    return
                # On any change to a .py file in the watched directory
                print(f"Detected change in {event.src_path}. Restarting MCP server...")
                self.manager.restart_with_debounce()
        return Handler(self)

    def start(self):
        # Ensure any existing process is fully stopped
        if self.process and self.process.poll() is None:
            print("Existing process still running, stopping it first...")
            self.stop()
            
        if self.process is None or self.process.poll() is not None:
            print(f"Starting MCP server with script: {self.main_script_path}")
            # Set PYTHONPATH for the subprocess to ensure it can find project modules
            env = os.environ.copy()
            project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
            env['PYTHONPATH'] = env.get('PYTHONPATH', '') + os.pathsep + project_root
            self.process = subprocess.Popen(['python', self.main_script_path], env=env)
            # Create a new observer if the current one has been stopped
            if not self.observer.is_alive():
                self.observer = Observer()
                self.event_handler = self._create_event_handler()
            # Watch the specified directory for changes
            self.observer.schedule(self.event_handler, self.watch_directory, recursive=True)
            if not self.observer.is_alive():
                self.observer.start()

    def stop(self):
        if self.process and self.process.poll() is None:
            print("Stopping MCP server...")
            self.process.terminate()
            try:
                # Wait for process to terminate, with timeout
                self.process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                print("Force killing MCP server process...")
                self.process.kill()
                self.process.wait()
        if self.observer.is_alive():
            self.observer.stop()
            # Don't join if we're being called from the observer thread itself
            import threading
            if threading.current_thread() != self.observer:
                self.observer.join()
        # Clear any scheduled events to avoid conflicts
        self.observer.unschedule_all()

    def restart_with_debounce(self):
        """Restart with debounce to prevent multiple rapid restarts"""
        current_time = time.time()
        if current_time - self.last_restart_time < self.restart_debounce_seconds:
            print(f"Restart request ignored due to debounce (last restart {current_time - self.last_restart_time:.1f}s ago)")
            return
        
        self.last_restart_time = current_time
        self.restart()
    
    def restart(self):
        # Use a separate thread to avoid deadlock when called from observer thread
        import threading
        def _restart():
            self.stop()
            # 等待更长时间确保端口完全释放
            time.sleep(5)
            
            # 检查端口是否真正释放
            import socket
            max_retries = 10
            for i in range(max_retries):
                try:
                    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                        s.bind(('localhost', 8000))
                        print(f"端口8000已释放，准备重启服务器")
                        break
                except OSError:
                    if i < max_retries - 1:
                        print(f"端口8000仍被占用，等待释放... ({i+1}/{max_retries})")
                        time.sleep(1)
                    else:
                        print("警告：端口可能仍被占用，强制重启")
            
            self.start()
        
        # Always use a separate thread to avoid blocking and potential deadlocks
        restart_thread = threading.Thread(target=_restart)
        restart_thread.daemon = True
        restart_thread.start()

    def is_running(self):
        return self.process is not None and self.process.poll() is None
    
    def create_backup(self):
        """创建当前脚本和function_registry.json的备份"""
        backup_created = False
        backup_info = {"script_backup": None, "registry_backup": None}
        
        if os.path.exists(self.main_script_path):
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            
            # 备份Python脚本
            script_backup_filename = f"mcp_server_backup_{timestamp}.py"
            script_backup_path = os.path.join(self.backup_dir, script_backup_filename)
            shutil.copy2(self.main_script_path, script_backup_path)
            backup_info["script_backup"] = script_backup_path
            print(f"Created script backup: {script_backup_path}")
            backup_created = True
            
            # 备份function_registry.json（如果存在）
            if os.path.exists(self.registry_path):
                registry_backup_filename = f"function_registry_backup_{timestamp}.json"
                registry_backup_path = os.path.join(self.backup_dir, registry_backup_filename)
                shutil.copy2(self.registry_path, registry_backup_path)
                backup_info["registry_backup"] = registry_backup_path
                print(f"Created registry backup: {registry_backup_path}")
        
        return backup_info if backup_created else None
    
    def health_check(self):
        """检查MCP服务器健康状态"""
        if not self.is_running():
            return False
        
        for attempt in range(self.health_check_retries):
            try:
                response = requests.get(self.health_check_url, timeout=5)
                if response.status_code == 200:
                    return True
            except requests.exceptions.RequestException:
                pass
            time.sleep(1)
        return False
    
    def start_with_validation(self):
        """启动服务器并验证健康状态"""
        # 尝试创建备份（如果文件存在的话）
        backup_info = self.create_backup()
        if backup_info:
            backup_files = []
            if backup_info["script_backup"]:
                backup_files.append(backup_info["script_backup"])
            if backup_info["registry_backup"]:
                backup_files.append(backup_info["registry_backup"])
            print(f"Created backup before starting: {', '.join(backup_files)}")
        else:
            print("No existing files to backup (first time startup)")
        
        # 启动服务器
        self.start()
        
        # 等待启动并检查健康状态
        start_time = time.time()
        while time.time() - start_time < self.startup_timeout:
            if self.health_check():
                print("MCP server started successfully and is healthy")
                self.last_known_good_script = self.main_script_path
                # 如果是第一次启动且没有备份，现在创建一个初始备份
                if not backup_info and os.path.exists(self.main_script_path):
                    initial_backup = self.create_backup()
                    if initial_backup:
                        backup_files = []
                        if initial_backup["script_backup"]:
                            backup_files.append(initial_backup["script_backup"])
                        if initial_backup["registry_backup"]:
                            backup_files.append(initial_backup["registry_backup"])
                        print(f"Created initial backup after successful startup: {', '.join(backup_files)}")
                return True
            time.sleep(2)
        
        print("MCP server failed health check after startup")
        return False
    
    def rollback_to_backup(self, max_attempts=None):
        """智能回退到可用的备份版本
        
        Args:
            max_attempts: 最大尝试次数，None表示尝试所有备份
        """
        # 检查备份目录是否存在
        if not os.path.exists(self.backup_dir):
            return {"success": False, "message": "备份目录不存在"}
        
        # 找到所有Python脚本备份文件
        script_backup_files = [f for f in os.listdir(self.backup_dir) if f.startswith("mcp_server_backup_") and f.endswith(".py")]
        if not script_backup_files:
            print("No script backup files found for rollback")
            return {"success": False, "message": "没有找到可用的脚本备份文件"}
        
        # 停止当前服务器
        self.stop()
        
        # 按时间戳排序，从最新开始尝试
        script_backup_files.sort(reverse=True)
        
        # 如果没有指定最大尝试次数，则尝试所有备份
        if max_attempts is None:
            max_attempts = len(script_backup_files)
        else:
            max_attempts = min(max_attempts, len(script_backup_files))
        
        # 保存当前文件作为临时备份
        temp_script_backup = None
        temp_registry_backup = None
        if os.path.exists(self.main_script_path):
            temp_script_backup = self.main_script_path + ".temp_backup"
            shutil.copy2(self.main_script_path, temp_script_backup)
        if os.path.exists(self.registry_path):
            temp_registry_backup = self.registry_path + ".temp_backup"
            shutil.copy2(self.registry_path, temp_registry_backup)
        
        # 逐个尝试备份文件
        for i in range(max_attempts):
            script_backup_file = script_backup_files[i]
            script_backup_path = os.path.join(self.backup_dir, script_backup_file)
            
            # 提取时间戳，查找对应的registry备份
            timestamp = script_backup_file.replace("mcp_server_backup_", "").replace(".py", "")
            registry_backup_file = f"function_registry_backup_{timestamp}.json"
            registry_backup_path = os.path.join(self.backup_dir, registry_backup_file)
            
            try:
                print(f"Attempting rollback to: {script_backup_file} (attempt {i+1}/{max_attempts})")
                
                # 恢复脚本备份
                shutil.copy2(script_backup_path, self.main_script_path)
                print(f"Restored script from: {script_backup_file}")
                
                # 恢复registry备份（如果存在）
                if os.path.exists(registry_backup_path):
                    shutil.copy2(registry_backup_path, self.registry_path)
                    print(f"Restored registry from: {registry_backup_file}")
                else:
                    print(f"No corresponding registry backup found: {registry_backup_file}")
                    # 如果没有对应的registry备份，创建空的registry
                    with open(self.registry_path, 'w', encoding='utf-8') as f:
                        f.write('{}')
                    print("Created empty registry file")
                
                # 尝试启动并验证
                validation_result = self.start_with_validation()
                if validation_result:
                    print(f"Successfully rolled back to: {script_backup_file}")
                    # 清理临时备份
                    if temp_script_backup and os.path.exists(temp_script_backup):
                        os.remove(temp_script_backup)
                    if temp_registry_backup and os.path.exists(temp_registry_backup):
                        os.remove(temp_registry_backup)
                    return {"success": True, "message": f"成功回退到备份: {script_backup_file}"}
                else:
                    print(f"Backup {script_backup_file} failed validation, trying next...")
                    self.stop()  # 停止失败的服务器
                    
            except Exception as e:
                print(f"Error testing backup {script_backup_file}: {str(e)}")
                self.stop()  # 确保停止任何可能启动的进程
                continue
        
        # 所有备份都失败了，尝试创建基础服务器
        print("All backups failed, creating fallback server...")
        try:
            self._create_basic_fallback_server()
            # 创建空的registry文件
            with open(self.registry_path, 'w', encoding='utf-8') as f:
                f.write('{}')
            print("Created empty registry file for fallback server")
            
            validation_result = self.start_with_validation()
            if validation_result:
                # 清理临时备份
                if temp_script_backup and os.path.exists(temp_script_backup):
                    os.remove(temp_script_backup)
                if temp_registry_backup and os.path.exists(temp_registry_backup):
                    os.remove(temp_registry_backup)
                return {"success": True, "message": "所有备份都失败，已创建基础服务器"}
        except Exception as e:
            print(f"Failed to create fallback server: {str(e)}")
        
        # 最后的尝试：恢复临时备份
        if temp_script_backup and os.path.exists(temp_script_backup):
            try:
                shutil.copy2(temp_script_backup, self.main_script_path)
                os.remove(temp_script_backup)
                print("Restored script from temp backup")
            except Exception as e:
                print(f"Failed to restore temp script backup: {str(e)}")
        
        if temp_registry_backup and os.path.exists(temp_registry_backup):
            try:
                shutil.copy2(temp_registry_backup, self.registry_path)
                os.remove(temp_registry_backup)
                print("Restored registry from temp backup")
            except Exception as e:
                print(f"Failed to restore temp registry backup: {str(e)}")
        
        return {"success": False, "message": f"尝试了 {max_attempts} 个备份版本，都无法正常启动"}
    
    def _create_basic_fallback_server(self):
        """创建最基础的MCP服务器脚本（无工具）"""
        basic_content = '''import sys
import os
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..'))

import logging
from datetime import datetime
from utils.MySQLUtil import MySQLUtil
from mcp.server.fastmcp import FastMCP
from fastmcp.exceptions import ToolError
from fastmcp.tools.tool import ToolAnnotations
from typing import Annotated
from pydantic import Field
import json
from starlette.requests import Request
from starlette.responses import JSONResponse

# Initialize FastMCP server with explicit configuration
mcp = FastMCP(
    name="xjcmcpServer",
    mask_error_details=True,
    host="0.0.0.0", port=8000
)

db = MySQLUtil()

logger = logging.getLogger("xjcmcp")
logging.basicConfig(level=logging.INFO)

# Health check endpoint
@mcp.custom_route("/health", methods=["GET"])
async def health_check(request: Request) -> JSONResponse:
    """Health check endpoint for monitoring server status"""
    try:
        # Test database connection
        db.test_connection()
        return JSONResponse({"status": "healthy", "timestamp": datetime.now().isoformat()})
    except Exception as e:
        logger.error(f"Health check failed: {e}")
        return JSONResponse({"status": "unhealthy", "error": str(e)}, status_code=500)

# No tools loaded - this is a basic fallback server

# Server execution block
if __name__ == "__main__":
    """Server entry point"""
    try:
        mcp.run(transport="sse")
    except KeyboardInterrupt:
        logger.info("Shutting down MCP server...")
'''
        
        with open(self.main_script_path, 'w', encoding='utf-8') as f:
            f.write(basic_content)
        print(f"Created basic fallback server at {self.main_script_path}")
    
    def get_server_status(self):
        """获取服务器详细状态"""
        script_backup_count = 0
        registry_backup_count = 0
        
        if os.path.exists(self.backup_dir):
            script_backup_count = len([f for f in os.listdir(self.backup_dir) if f.startswith("mcp_server_backup_") and f.endswith(".py")])
            registry_backup_count = len([f for f in os.listdir(self.backup_dir) if f.startswith("function_registry_backup_") and f.endswith(".json")])
        
        status = {
            "is_running": self.is_running(),
            "is_healthy": False,
            "last_restart_time": self.last_restart_time,
            "script_backup_count": script_backup_count,
            "registry_backup_count": registry_backup_count,
            "backup_count": script_backup_count  # 保持向后兼容
        }
        
        if status["is_running"]:
            status["is_healthy"] = self.health_check()
        
        return status